#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
加密密钥生成和管理工具
使用方法:
    python crypto_key_manager.py generate --algorithm random --length 32
    python crypto_key_manager.py generate --algorithm sha256
    python crypto_key_manager.py generate --algorithm rsa --key-size 2048
    python crypto_key_manager.py generate --algorithm aes --key-size 256
    python crypto_key_manager.py test --algorithm aes --key your-key-here
"""

import argparse
import sys
import os
import json
from app.utils.crypto_utils import CryptoUtils


def generate_key(args):
    """生成密钥"""
    try:
        if args.algorithm == 'md5':
            data = args.data or "default-data"
            key = CryptoUtils.generate_md5_hash(data)
            print(f"MD5哈希: {key}")
            
        elif args.algorithm == 'sha1':
            data = args.data or "default-data"
            key = CryptoUtils.generate_sha1_hash(data)
            print(f"SHA1哈希: {key}")
            
        elif args.algorithm == 'sha256':
            data = args.data or "default-data"
            key = CryptoUtils.generate_sha256_hash(data)
            print(f"SHA256哈希: {key}")
            
        elif args.algorithm == 'sha512':
            data = args.data or "default-data"
            key = CryptoUtils.generate_sha512_hash(data)
            print(f"SHA512哈希: {key}")
            
        elif args.algorithm == 'random':
            length = args.length or 32
            key = CryptoUtils.generate_random_secret(length)
            print(f"随机密钥({length}位): {key}")
            
        elif args.algorithm == 'hex':
            length = args.length or 32
            key = CryptoUtils.generate_hex_secret(length)
            print(f"十六进制密钥({length}位): {key}")
            
        elif args.algorithm == 'base64':
            length = args.length or 32
            key = CryptoUtils.generate_base64_secret(length)
            print(f"Base64密钥({length}位): {key}")
            
        elif args.algorithm == 'rsa':
            key_size = args.key_size or 2048
            private_key, public_key = CryptoUtils.generate_rsa_key_pair(key_size)
            print(f"RSA私钥({key_size}位):\n{private_key}\n")
            print(f"RSA公钥({key_size}位):\n{public_key}")
            
            # 如果指定了输出文件，保存密钥对
            if args.output:
                key_pair = {
                    "private_key": private_key,
                    "public_key": public_key
                }
                with open(args.output, 'w') as f:
                    json.dump(key_pair, f, indent=2)
                print(f"\nRSA密钥对已保存到: {args.output}")
                
        elif args.algorithm == 'aes':
            key_size = args.key_size or 256
            key = CryptoUtils.generate_aes_key(key_size)
            print(f"AES密钥({key_size}位): {key}")
            
        elif args.algorithm == 'des':
            key = CryptoUtils.generate_des_key()
            print(f"DES密钥: {key}")
            
        elif args.algorithm == 'jwt':
            algorithm = args.jwt_algorithm or 'random'
            length = args.length or 32
            key = CryptoUtils.generate_jwt_secret(algorithm, length)
            print(f"JWT密钥({algorithm}): {key}")
            
        elif args.algorithm == 'hmac':
            data = args.data or "default-data"
            key = args.key or "default-key"
            hash_algorithm = args.hash_algorithm or 'sha256'
            hmac_hash = CryptoUtils.generate_hmac_hash(data, key, hash_algorithm)
            print(f"HMAC-{hash_algorithm}哈希: {hmac_hash}")
            
        elif args.algorithm == 'pbkdf2':
            password = args.password or "default-password"
            salt = args.salt
            iterations = args.iterations or 100000
            key_length = args.key_length or 32
            algorithm = args.hash_algorithm or 'sha256'
            
            hash_value, salt_used = CryptoUtils.generate_pbkdf2_hash(
                password, salt, iterations, key_length, algorithm
            )
            print(f"PBKDF2-{algorithm}哈希: {hash_value}")
            print(f"使用的盐值: {salt_used}")
            
        else:
            print(f"不支持的算法: {args.algorithm}")
            return 1
            
        return 0
        
    except Exception as e:
        print(f"生成密钥时出错: {str(e)}")
        return 1


def test_crypto(args):
    """测试加密/解密功能"""
    try:
        if args.algorithm == 'aes':
            key = args.key
            if not key:
                print("AES测试需要提供密钥 (--key)")
                return 1
                
            plaintext = args.data or "这是一段测试文本"
            print(f"原始文本: {plaintext}")
            
            # 加密
            encrypted = CryptoUtils.aes_encrypt(plaintext, key)
            print(f"加密后: {encrypted}")
            
            # 解密
            decrypted = CryptoUtils.aes_decrypt(encrypted, key)
            print(f"解密后: {decrypted}")
            
            # 验证
            if plaintext == decrypted:
                print("✓ AES加密/解密测试成功")
                return 0
            else:
                print("✗ AES加密/解密测试失败")
                return 1
                
        elif args.algorithm == 'rsa':
            private_key = args.private_key
            public_key = args.public_key
            
            if not private_key or not public_key:
                print("RSA测试需要提供私钥 (--private-key) 和公钥 (--public-key)")
                return 1
                
            plaintext = args.data or "这是一段测试文本"
            print(f"原始文本: {plaintext}")
            
            # 加密
            encrypted = CryptoUtils.rsa_encrypt(plaintext, public_key)
            print(f"加密后: {encrypted}")
            
            # 解密
            decrypted = CryptoUtils.rsa_decrypt(encrypted, private_key)
            print(f"解密后: {decrypted}")
            
            # 验证
            if plaintext == decrypted:
                print("✓ RSA加密/解密测试成功")
                return 0
            else:
                print("✗ RSA加密/解密测试失败")
                return 1
                
        else:
            print(f"不支持的测试算法: {args.algorithm}")
            return 1
            
    except Exception as e:
        print(f"测试加密/解密时出错: {str(e)}")
        return 1


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='加密密钥生成和管理工具')
    subparsers = parser.add_subparsers(dest='command', help='可用命令')
    
    # 生成密钥命令
    generate_parser = subparsers.add_parser('generate', help='生成密钥')
    generate_parser.add_argument('--algorithm', required=True, 
                                choices=['md5', 'sha1', 'sha256', 'sha512', 'random', 
                                        'hex', 'base64', 'rsa', 'aes', 'des', 'jwt', 
                                        'hmac', 'pbkdf2'],
                                help='加密算法')
    generate_parser.add_argument('--length', type=int, help='密钥长度')
    generate_parser.add_argument('--key-size', type=int, help='密钥长度(适用于RSA/AES)')
    generate_parser.add_argument('--data', help='要哈希的数据')
    generate_parser.add_argument('--key', help='密钥')
    generate_parser.add_argument('--password', help='密码')
    generate_parser.add_argument('--salt', help='盐值')
    generate_parser.add_argument('--iterations', type=int, help='迭代次数')
    generate_parser.add_argument('--key-length', type=int, help='密钥长度')
    generate_parser.add_argument('--hash-algorithm', default='sha256', 
                                help='哈希算法(适用于HMAC/PBKDF2)')
    generate_parser.add_argument('--jwt-algorithm', default='random', 
                                help='JWT密钥生成算法')
    generate_parser.add_argument('--output', help='输出文件路径(适用于RSA)')
    generate_parser.set_defaults(func=generate_key)
    
    # 测试加密/解密命令
    test_parser = subparsers.add_parser('test', help='测试加密/解密功能')
    test_parser.add_argument('--algorithm', required=True, 
                            choices=['aes', 'rsa'],
                            help='加密算法')
    test_parser.add_argument('--key', help='密钥(适用于AES)')
    test_parser.add_argument('--private-key', help='私钥(适用于RSA)')
    test_parser.add_argument('--public-key', help='公钥(适用于RSA)')
    test_parser.add_argument('--data', help='要加密的数据')
    test_parser.set_defaults(func=test_crypto)
    
    # 解析参数
    args = parser.parse_args()
    
    # 执行命令
    if hasattr(args, 'func'):
        return args.func(args)
    else:
        parser.print_help()
        return 1


if __name__ == '__main__':
    sys.exit(main())